from tensorflow.python import pywrap_tensorflow
import os

bert_dir = '/home/ahashi_syuu/Documents/BERT/uncased_L-12_H-768_A-12/uncased_L-12_H-768_A-12/'
checkpoint_path = bert_dir + 'bert_model.ckpt'
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_map = reader.get_variable_to_shape_map()
for key in var_map:
    print('tensor name: ', key, '\t', reader.get_tensor(key).shape)

print('end')

